import asyncio, json
from loguru import logger

from core.command_manager import CommandEvent


class TCPClient:
    def __init__(self, host="127.0.0.1", port=9000, state=None):
        self.host = host
        self.port = port
        self.reader=None
        self.writer=None
        self.handlers = {}
        self.state = state
        self.heartbeat = 5  # 心跳间隔改为5秒，提高状态同步速度
        self._reconnect_delay = 3

    def register_handler(self, cmd, coro):
        self.handlers[cmd]=coro

    async def connect_loop(self):
        # ✅ 检查全局关闭事件（从main模块导入）
        from __main__ import shutdown_event
        
        while not shutdown_event.is_set():
            try:
                logger.info(f"正在连接... {self.host}:{self.port}")
                if self.state:
                    self.state.update(connected=False)
                    
                # ✅ 解析host（如果是域名，解析为IP）
                import socket
                try:
                    # 尝试解析host（如果是IP地址，直接使用；如果是域名，解析为IP）
                    resolved_ip = socket.gethostbyname(self.host)
                    if resolved_ip != self.host:
                        logger.info(f"✅ 域名解析: {self.host} -> {resolved_ip}")
                    # ✅ 保存控制端IP到state（用于文件接收服务器的IP白名单）
                    if self.state:
                        self.state.update(control_server_ip=resolved_ip)
                        logger.info(f"✅ 控制端IP已保存到state: {resolved_ip}")
                except Exception as e:
                    logger.warning(f"⚠️ 解析控制端地址失败: {e}，使用原始host: {self.host}")
                    # 如果解析失败，假设host就是IP
                    if self.state:
                        self.state.update(control_server_ip=self.host)
                
                # ✅ 添加连接超时（10秒），避免防火墙未开放端口时长时间等待
                try:
                    self.reader, self.writer = await asyncio.wait_for(
                        asyncio.open_connection(self.host, self.port),
                        timeout=10.0
                    )
                    logger.info(f"✅ 连接成功: {self.host}:{self.port}")
                except asyncio.TimeoutError:
                    raise Exception(f"连接超时: {self.host}:{self.port} (10秒内未响应，请检查：1. 控制端是否启动 2. 防火墙是否开放{self.port}端口 3. 网络是否连通)")
                except ConnectionRefusedError:
                    raise Exception(f"连接被拒绝: {self.host}:{self.port} (控制端可能未启动或防火墙未开放{self.port}端口)")
                except Exception as e:
                    raise Exception(f"连接失败: {self.host}:{self.port} - {str(e)}")
                
                if self.state:
                    self.state.update(connected=True, online=True)
                
                # 立即发送一次状态消息，让控制端第一时间知道客户端上线（包含cookie）
                logger.info("📤 准备发送初始状态消息...")
                try:
                    await self.send_status(include_cookie=True)
                    logger.info("✅ 已发送初始状态消息（包含cookie）")
                except Exception as e:
                    logger.error(f"❌ 发送初始状态消息失败: {e}", exc_info=True)
                
                # 启动心跳循环
                logger.info("💓 启动心跳循环...")
                heartbeat_task = asyncio.create_task(self._heartbeat_loop())
                
                # 读取数据循环
                logger.info("📡 启动读取数据循环...")
                await self._reader_loop()
                
                # 取消心跳任务
                heartbeat_task.cancel()
                
                # ✅ 更新状态为离线
                if self.state:
                    self.state.update(connected=False, online=False)
                
                # ✅ 检查关闭事件，如果已设置则退出
                if shutdown_event.is_set():
                    logger.info("收到关闭信号，停止连接")
                    break
                
                # ✅ 连接断开，准备重连
                logger.warning(f"连接已断开，{self._reconnect_delay}秒后自动重连...")
                
            except Exception as e:
                logger.error(f"连接错误: {e}")
                # ✅ 连接失败时清空连接对象
                if self.writer:
                    try:
                        self.writer.close()
                        await self.writer.wait_closed()
                    except:
                        pass
                self.writer = None
                self.reader = None
                
                if self.state:
                    self.state.update(connected=False, online=False)
                
                # ✅ 检查关闭事件，如果已设置则退出
                if shutdown_event.is_set():
                    logger.info("收到关闭信号，停止连接")
                    break
                    
                logger.warning(f"连接失败，{self._reconnect_delay}秒后自动重连...")
            finally:
                if not shutdown_event.is_set():
                    await asyncio.sleep(self._reconnect_delay)
                else:
                    break


    async def _reader_loop(self):
        logger.info("开始读取数据")
        from __main__ import shutdown_event
        read_timeout_count = 0
        
        while not shutdown_event.is_set():
            try:
                # ✅ 读取一行数据，带超时以便检查 shutdown_event
                try:
                    line = await asyncio.wait_for(self.reader.readline(), timeout=1.0)
                except asyncio.TimeoutError:
                    # 超时是正常的，继续循环检查 shutdown_event
                    read_timeout_count += 1
                    # ✅ 每60次超时（约1分钟）输出一次日志，确认读取循环还在运行
                    if read_timeout_count % 60 == 0:
                        logger.info(f"📡 读取循环运行中... (已等待 {read_timeout_count} 秒，连接状态: {'✅ 已连接' if self.writer else '❌ 未连接'})")
                    continue
                
                if not line:
                    logger.warning("⚠️ 连接已被控制端断开（读取到EOF，通常表示控制端主动关闭连接）")
                    logger.debug("可能原因：1. 控制端重启/更新 2. 心跳超时检测 3. 网络问题 4. 控制端资源清理")
                    break
                
                # ✅ 立即记录收到的原始数据（用于调试）
                text = line.decode().strip()
                # ✅ 使用INFO级别，确保能看到所有收到的数据
                logger.info(f"📥 收到原始数据: {text[:500]}")  # 限制长度避免日志过长
                
                msg = CommandEvent.parse(text)
                if msg is None:
                    logger.warning(f"⚠️ 无法解析数据: {text[:200]}")
                    continue

                # ✅ 记录所有收到的命令（除了chat_open和status）
                event_type = msg.event.value
                if event_type == "chat_open":
                    # chat_open命令太频繁，只记录DEBUG级别
                    logger.debug(f"📥 收到命令: chat_open")
                elif event_type == "status":
                    # 状态消息（心跳）只在DEBUG级别记录
                    logger.debug(f"📥 收到状态消息（心跳）")
                else:
                    # ✅ 其他所有命令都记录INFO级别，确保能看到
                    logger.info(f"📥 收到命令: {event_type}, data={msg.data}")

                asyncio.create_task(self._handle(msg))
            except ConnectionResetError:
                logger.warning("连接被重置")
                break
            except Exception as e:
                logger.error(f"❌ 读取错误: {e}", exc_info=True)
                break
        
        # ✅ 快速关闭连接
        try:
            if self.writer:
                self.writer.close()
                try:
                    await asyncio.wait_for(self.writer.wait_closed(), timeout=0.5)
                except:
                    pass
        except:
            pass


    async def _handle(self, msg):
        cmd = msg.event
        handler = self.handlers.get(cmd)
        
        if handler:
            try:
                # chat_open命令太多，不记录日志
                if cmd.value != "chat_open":
                    logger.debug(f"执行命令处理器: {cmd.value}")
                res = handler(msg)
                if asyncio.iscoroutine(res):
                    out = await res
                    if isinstance(out, dict):
                        # ✅ 精简日志：chat_open命令不记录日志
                        if cmd.value == "chat_open":
                            pass  # 不记录日志
                        elif cmd.value in ["chat_send", "cancel_order"]:
                            if out.get("success"):
                                logger.success(f"✅ {cmd.value}成功: {out.get('data', {})}")
                            else:
                                logger.error(f"❌ {cmd.value}失败: {out.get('msg')}")
                    await self.send(out)
                elif isinstance(res, dict):
                    # ✅ 记录重要命令的执行结果
                    if cmd.value in ["chat_open", "chat_send", "cancel_order"]:
                        if res.get("success"):
                            logger.success(f"✅ {cmd.value}成功: {res.get('data', {})}")
                        else:
                            logger.error(f"❌ {cmd.value}失败: {res.get('msg')}")
                    await self.send(res)
            except Exception as e:
                logger.error(f"❌ 执行命令处理器 {cmd.value} 失败: {e}", exc_info=True)
                from core.command_manager import CommandRes
                await self.send(CommandRes.Fail(msg=str(e)).model_dump())
        else:
            logger.warning(f"⚠️ 事件{cmd.value}不存在对应的处理器，已注册的handlers={[h.value for h in self.handlers.keys()]}")


    async def send(self, payload: dict):
        # ✅ 检查连接是否已建立
        if self.writer is None:
            logger.debug(f"⚠️ 连接未建立，无法发送消息: {payload.get('event', 'unknown')}")
            return
        
        try:
            data = (json.dumps(payload, ensure_ascii=False) + "\n").encode()
            self.writer.write(data)
            await self.writer.drain()
        except Exception as e:
            logger.error(f"发送错误: {e}")
            # ✅ 连接可能已断开，清空writer和reader
            self.writer = None
            self.reader = None


    async def _heartbeat_loop(self):
        from __main__ import shutdown_event
        heartbeat_count = 0
        while not shutdown_event.is_set():
            try:
                # ✅ 等待 shutdown_event 或超时，超时后发送心跳
                # 使用 wait_for + shutdown_event.wait() 可以被立即中断
                try:
                    await asyncio.wait_for(shutdown_event.wait(), timeout=self.heartbeat)
                    # shutdown_event 被设置了，退出循环
                    break
                except asyncio.TimeoutError:
                    # 超时了，说明还没有关闭信号，发送心跳
                    if not shutdown_event.is_set():
                        heartbeat_count += 1
                        # ✅ 每10次心跳输出一次日志（减少日志量，但能看到心跳在工作）
                        if heartbeat_count % 10 == 0:
                            logger.info(f"💓 心跳 #{heartbeat_count} (连接状态: {'✅ 已连接' if self.writer else '❌ 未连接'})")
                        await self.send_status()
            except asyncio.CancelledError:
                logger.debug("心跳循环被取消")
                break
            except Exception as e:
                logger.error(f"心跳循环异常: {e}", exc_info=True)
                # 出错后等待一下再重试
                await asyncio.sleep(1)


    async def send_status(self, include_cookie=False):
        """发送状态信息到控制端
        
        Args:
            include_cookie: 是否包含cookie（默认False，减少心跳包大小）
        """
        from core.command_manager import CommandRes
        if self.state:
            self.state.mark_heartbeat()
            payload = CommandRes.Status().Success(data=self.state.to_dict(include_cookie=include_cookie))
        else:
            payload = CommandRes.Status().Success(data={"online": True})
        
        # ✅ 检查连接状态
        if self.writer is None:
            logger.warning("⚠️ 尝试发送心跳，但连接未建立")
            return
        
        await self.send(payload)
    
    async def send_cookie_update(self):
        """单独发送cookie更新（使用update_cookie事件）"""
        from core.command_manager import CommandRes
        if self.state and self.state.get('cookie'):
            payload = {
                "event": "update_cookie",
                "direction": "res",
                "success": True,
                "code": 0,
                "msg": "",
                "data": {
                    "client_id": self.state.get('client_id'),
                    "cookie": self.state.get('cookie')
                }
            }
            await self.send(payload)
            logger.info("✅ 已发送cookie更新")
